import colorsys

from utils.dataset import get_classes

def visual_color(classes_path):
    """
    自动为不同的类别分配颜色
    """
    class_names, num_classes = get_classes(classes_path)
    #---------------------------------------------------#
    #   画框设置不同的颜色
    #---------------------------------------------------#
    hsv_tuples = [(x / num_classes, 1., 1.) for x in range(num_classes)]
    colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
    colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), colors))
    
    colors = dict(zip(class_names, colors))
    return colors
    
if __name__ == '__main__':
    colors = visual_color()
    print(colors)